from datetime import datetime
from pymongo import ASCENDING, IndexModel
from pymongo.collection import Collection

from ..interface import Job, JobDoneUpdateModel, JobRunTimesUpdateModel, JobStore


class MongoStore(JobStore):
    """Mongodb任务存储实现"""
    collection: Collection

    def __init__(self, collection):
        self.collection = collection
        idx1 = IndexModel([('func_name', ASCENDING), ('done', ASCENDING)])
        idx2 = IndexModel([('job_id', ASCENDING)], unique=True)
        self.collection.create_indexes([idx1, idx2])

    def get_by_job_id(self, job_id: str) -> Job:
        """通过任务ID获取任务记录"""
        document = self.collection.find_one({'job_id': job_id})
        return Job(**document)

    def add(self, job: Job):
        now = datetime.now()
        job.update_time = now
        job.create_time = now
        result = self.collection.insert_one(job.dict())
        return result.inserted_id

    def _update_one(self, job_id: str, data: dict) -> bool:
        data['update_time'] = datetime.now()
        result = self.collection.update_one(
            {'job_id': job_id},
            {'$set': data}
        )
        return result.modified_count == 1
    
    def update_run_times(self, job_id: str, m: JobRunTimesUpdateModel) -> bool:
        return self._update_one(job_id, m.dict())

    def done(self, job_id: str, m: JobDoneUpdateModel) -> bool:
        return self._update_one(job_id, m.dict())